﻿module FES2DT3NC01_m
use AFES_m
use Mesh2DT3CR_m
use Mesh2DT0_m
use Mesh2DT3_m
use FES2DT3CR1_m
type, extends(AFES) :: FES2DT3NC01
    type(Mesh2DT3CR), pointer :: meshVel
    type(Mesh2DT0), pointer :: meshPre
contains
    procedure :: initFES => initFES2DT3NC01
    procedure :: valueFromFunc2DT3CR
    procedure :: valueFromFunc2DT0
    procedure :: assembleStokes2DT3NC01_Const
    generic :: assembleStokes => assembleStokes2DT3NC01_Const
    
    procedure :: applyBCZero2DT3NC01_MR
    generic :: applyBCZero => applyBCZero2DT3NC01_MR
    
    procedure :: applyBC2DT3NC01_function
    procedure :: applyBC2DT3NC01_const
    generic :: applyBC => applyBC2DT3NC01_function, applyBC2DT3NC01_const
    
    procedure :: extractUVP => extractUVP2DT3NC01
end type
contains

subroutine initFES2DT3NC01(this, meshCR, meshT0)
    implicit none
    class(FES2DT3NC01), intent(out) :: this
    type(Mesh2DT3CR), intent(in), target :: meshCR
    type(Mesh2DT0), intent(in), target :: meshT0
    
    this.meshVel => meshCR
    this.meshPre => meshT0
    this.dim = 2
    this.nDof = 2 * meshCR.nv + meshT0.nv + 1
    this.nVertices = meshCR.nv + meshT0.nv
    this.nElement = meshCR.nt
    this.nBoundary = meshCR.nb
    this.nPerElement = 4
    this.nPerBoundary = 1
    this.dofPerNode = 3
    this.dofPerElement = 7
    this.dofPerBoundary = 2
    this.nBand = 2 * meshCR.nv + meshT0.nv + 1
end subroutine

subroutine valueFromFunc2DT3CR(this, func, u)
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    real*8, external :: func
    real*8, intent(out) :: u(this.meshVel.nv)
    
    integer :: i
    type(Mesh2DT3CR), pointer :: mesh2D
    mesh2D => this%meshVel
    do i = 1, mesh2D%nv
        u(i) = func(mesh2D%nodes(i)%xy)
    end do
end subroutine

subroutine valueFromFunc2DT0(this, func, u)
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    real*8, external :: func
    real*8, intent(out) :: u(this.meshPre.nv)
    
    integer :: i
    type(Mesh2DT0), pointer :: mesh2D
    mesh2D => this%meshPre
    do i = 1, mesh2D%nv
        u(i) = func(mesh2D%nodes(i)%xy)
    end do
end subroutine


subroutine assembleStokes2DT3NC01_Const(this, nu, A)
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    real*8, intent(in) :: nu
    type(SMatrix), intent(out) :: A

    integer :: i, j, nv
    real*8 :: xy(3, 2), nnu(3), dxP(3), dyP(3)
    real*8, allocatable :: areaT(:)
    integer, allocatable :: areaIdx(:)
    type(mesh2DT3CR), pointer :: mesh2D
    integer :: idx(3), nEdge
    real*8 :: K(3, 3)
    nnu = nu
    mesh2D => this%meshVel
    nv = mesh2D%nv
    allocate(areaT(mesh2D%nt))
    call this.meshVel.meshBG.area(areaT)
    allocate(areaIdx(mesh2D%nt))
    do i = 1, mesh2D%nt
        idx = mesh2D%elements(i)%idx
        xy(:, 1) = mesh2D%nodes(idx)%xy(1)
        xy(:, 2) = mesh2D%nodes(idx)%xy(2)
        call elementHeatStiff2DT3CR1(xy, nnu, COMMON_2D, K)
        call elementPressPartial(xy, dxP, dyP)
        idx = mesh2D%elements(i)%idx
        call assembleMatrix(A, K, idx)
        call assembleMatrixCol(A, dxP, idx, i + 2 * nv)
        call assembleMatrixRow(A, dxP, i + 2 * nv, idx)
        idx = idx + nv
        call assembleMatrix(A, K, idx)
        call assembleMatrixCol(A, dyP, idx, i + 2 * nv)       
        call assembleMatrixRow(A, dyP, i + 2 * nv, idx)
        areaIdx(i) = 2 * nv + i
    end do
    call assembleMatrixCol(A, areaT, areaIdx, this%nDof)
    call assembleMatrixRow(A, areaT, this%nDof, areaIdx)
    deallocate(areaT, areaIdx)
end subroutine

subroutine elementPressPartial(xy, dxP, dyP)
    implicit none
    real*8, intent(in) :: xy(3, 2)
    real*8, intent(out) :: dxP(3), dyP(3)
    
    real*8 :: dx(3), dy(3), x(3), y(3), area
    x = xy(:, 1); y = xy(:, 2)
    call gradientCR(x, y, dx, dy, area)
    dxP = -dx * area
    dyP = -dy * area
end subroutine

subroutine applyBCZero2DT3NC01_MR(this, A, RHS, label, direct)
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: RHS(:)
    integer, intent(in) :: label(:)
    integer, optional, intent(in) :: direct
    
    type(LinkedSet) :: dirichletBC
    call this.meshVel.extractBoundaryNodes(dirichletBC, label)
    if (present(direct)) then
        call this.applyBCZero1_MR(A, RHS, dirichletBC, direct)
    else
        call this.applyBCZero1_MR(A, RHS, dirichletBC)
    end if
	call dirichletBC.empty()
end subroutine

subroutine applyBC2DT3NC01_function(this, A, RHS, label, u0, v0, time, tp, direct)
    use Set_m
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: RHS(:)
    integer, intent(in) :: label(:), tp
    real*8, external :: u0, v0
    real*8, intent(in) :: time
    integer, intent(in), optional :: direct
    
    real*8, allocatable :: uu0(:)
    integer, allocatable :: BCDOF(:)
    integer :: i
    real*8 :: xy(2)
    type(Mesh2DT3CR), pointer :: meshV
    type(LinkedSet) :: BSet
    type(ele), pointer :: p
    meshV => this%meshVel
    
    call meshV.extractBoundaryNodes(BSet, label)
    
    if (present(direct)) then
        allocate(uu0(BSet%size))
        allocate(BCDOF(BSet%size))
        i = 0
        p => BSet%head
        if (direct == directX) then
            do while(associated(p))
                i = i + 1
                xy = meshV%nodes(p%n)%xy
                uu0(i) = u0(xy, time)
                BCDOF(i) = p%n
                p => p%next
            end do
        else if (direct == directY) then
            do while(associated(p))
                i = i + 1
                xy = meshV%nodes(p%n)%xy
                uu0(i) = v0(xy, time)
                BCDOF(i) = p%n + meshV%nv
                p => p%next
            end do
        end if
        call this.applyBC_DOF(A, RHS, BCDOF, uu0, tp)
    else
        allocate(uu0(2 * BSet%size))
        allocate(BCDOF(2 * BSet%size))
        i = 0
        p => BSet%head
        do while(associated(p))
            i = i + 1
            xy = meshV%nodes(p%n)%xy
            BCDOF(2 * i - 1) = p%n
            BCDOF(2 * i) = p%n + meshV%nv
            uu0(2 * i - 1) = u0(xy, time)
            uu0(2 * i) = v0(xy, time)
            p => p%next
        end do
        call this.applyBC_DOF(A, RHS, BCDOF, uu0, tp)
    end if
end subroutine

subroutine applyBC2DT3NC01_const(this, A, RHS, label, u0, v0, tp, direct)
    use Set_m
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: RHS(:)
    integer, intent(in) :: label(:), tp
    real*8, intent(in) :: u0, v0
    integer, intent(in), optional :: direct
    
    real*8, allocatable :: uu0(:)
    integer, allocatable :: BCDOF(:)
    integer :: i
    type(Mesh2DT3CR), pointer :: meshV
    type(LinkedSet) :: BSet
    type(ele), pointer :: p
    meshV => this%meshVel
    
    call meshV.extractBoundaryNodes(BSet, label)
    
    if (present(direct)) then
        allocate(uu0(BSet%size))
        allocate(BCDOF(BSet%size))
        i = 0
        p => BSet%head
        if (direct == directX) then
            uu0 = u0
            do while(associated(p))
                i = i + 1
                BCDOF(i) = p%n
                p => p%next
            end do
        else if (direct == directY) then
            uu0 = v0
            do while(associated(p))
                i = i + 1
                BCDOF(i) = p%n + meshV%nv
                p => p%next
            end do
        end if
        call this.applyBC_DOF(A, RHS, BCDOF, uu0, tp)
    else
        allocate(uu0(2 * BSet%size))
        allocate(BCDOF(2 * BSet%size))
        i = 0
        p => BSet%head
        do while(associated(p))
            i = i + 1
            BCDOF(2 * i - 1) = p%n
            BCDOF(2 * i) = p%n + meshV%nv
            uu0(2 * i - 1) = u0
            uu0(2 * i) = v0
            p => p%next
        end do
        call this.applyBC_DOF(A, RHS, BCDOF, uu0, tp)
    end if
end subroutine


subroutine extractUVP2DT3NC01(this, x, u, v, p)
    implicit none
    class(FES2DT3NC01), intent(in) :: this
    real*8, intent(in) :: x(this%nDof)
    real*8, intent(out) :: u(:)
    real*8, intent(out) :: v(:)
    real*8, intent(out) :: p(:)
    
    type(Mesh2DT3CR), pointer :: meshV
    type(Mesh2DT0), pointer :: meshP
    meshV => this%meshVel
    meshP => this%meshPre
    u = x(1:meshV%nv)
    v = x(meshV%nv+1:2 * meshV%nv)
    p = x(2 * meshV%nv + 1 : this%nDof - 1)
end subroutine


end module